package com.zhentao.util;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树形结构构建工具类
 * @param <T> 节点类型
 * @param <K> 节点ID类型（如Long、String）
 */
public class TreeBuilder<T, K> {

    /**
     * 将扁平列表转换为树形结构
     * @param list 原始数据列表
     * @param idGetter 获取节点ID的方法（如 Menu::getId）
     * @param parentIdGetter 获取父节点ID的方法（如 Menu::getParentId）
     * @param childrenSetter 设置子节点的方法（如 Menu::setChildren）
     * @param rootParentId 根节点的父ID值（通常为0或null）
     * @return 树形结构列表
     */
    public List<T> build(
            List<T> list,
            Function<T, K> idGetter,
            Function<T, K> parentIdGetter,
            ChildrenSetter<T> childrenSetter,
            K rootParentId) {

        // 按父ID分组，加速子节点查找
        Map<K, List<T>> parentIdToChildren = list.stream()
                .filter(node -> !Objects.equals(parentIdGetter.apply(node), rootParentId))
                .collect(Collectors.groupingBy(parentIdGetter));

        // 构建树
        return list.stream()
                .filter(node -> Objects.equals(parentIdGetter.apply(node), rootParentId))
                .peek(node -> setChildren(node, parentIdToChildren, idGetter, childrenSetter))
                .collect(Collectors.toList());
    }

    // 递归设置子节点
    private void setChildren(
            T node,
            Map<K, List<T>> parentIdToChildren,
            Function<T, K> idGetter,
            ChildrenSetter<T> childrenSetter) {

        K nodeId = idGetter.apply(node);
        List<T> children = parentIdToChildren.getOrDefault(nodeId, Collections.emptyList());
        if (!children.isEmpty()) {
            children.forEach(child -> setChildren(child, parentIdToChildren, idGetter, childrenSetter));
            childrenSetter.setChildren(node, children);
        }
    }

    /**
     * 子节点设置器（函数式接口）
     */
    @FunctionalInterface
    public interface ChildrenSetter<T> {
        void setChildren(T parent, List<T> children);
    }
}